from sklearn import datasets
import matplotlib.pyplot as plt
import numpy as np

# 获取数据集并进行探索
iris = datasets.load_iris()
irisFeatures = iris["data"]
irisFeaturesName = iris["feature_names"]
irisLabels = iris["target"]


def norm2(x):
    # 求2范数的平方值
    return np.sum(x * x)



def drawSL(cnt, J_oldJ):
    # cnt为x轴坐标，J_oldJ为y轴坐标
    plt.plot(range(cnt), J_oldJ, "r-")
    plt.xlabel("iteration")
    plt.ylabel("abs(J-oldJ)")
    plt.show()


class KMeans(object):
    def __init__(self, k: int, n: int):  # k:聚类的数目，n:数据维度
        self.K = k
        self.N = n
        self.u = np.zeros((k, n))
        self.C = [[] for i in range(k)]
        # u[i]：第i个聚类中心，C[i]：第i个类别所包含的点

    def fit(self, data: np.ndarray):
        # data:每一行是一个样本
        self.select_u0(data)
        # 输出聚类中心初始值
        print("初始聚类中心：")
        print(self.u)

        cnt = 0
        J_oldJ = []

        # 聚类中心初始化
        J = 0
        oldJ = 100
        while abs(J - oldJ) > 0.001:
            oldJ = J
            J = 0
            cnt += 1

            self.C = [[] for i in range(self.K)]
            for x in data:
                nor = [norm2(self.u[i] - x) \
                       for i in range(self.K)]
                J += np.min(nor)
                self.C[np.argmin(nor)].append(x)
            self.u = [np.mean(np.array(self.C[i]), axis=0) \
                      for i in range(self.K)]

            J_oldJ.append(float(abs(J - oldJ)))

        drawSL(cnt, J_oldJ)

    def select_u0(self, data: np.ndarray):
        for j in range(self.N):
            # 得到该列数据的最小值,最大值

            minJ = np.min(data[:, j])

            maxJ = np.max(data[:, j])

            rangeJ = float(maxJ - minJ)
            # 聚类中心的第j维数据值随机为位于(最小值，最大值)内
            self.u[:, j] = minJ + rangeJ * np.random.rand(self.K)


#  k为聚类的数目，即图中取几个聚类的点
model = KMeans(1, 4)
# k=3，n=4
model.fit(irisFeatures)

# x = np.array(model.C[0])
# plt.scatter(x[:, 0], x[:, 1], c="red", marker='o', label='cluster1')
# x = np.array(model.C[1])
# plt.scatter(x[:, 0], x[:, 1], c="green", marker='*', label='cluster2')
# x = np.array(model.C[2])
# plt.scatter(x[:, 0], x[:, 1], c="blue", marker='+', label='cluster3')
# x = np.array(model.C[3])
# plt.scatter(x[:, 0], x[:, 1], c="purple", marker='1', label='cluster4')
# x = np.array(model.C[4])
# plt.scatter(x[:, 0], x[:, 1], c="magenta", marker='2', label='cluster5')
# x = np.array(model.C[5])
# plt.scatter(x[:, 0], x[:, 1], c="cyan", marker='3', label='cluster6')
# x = np.array(model.C[6])
# plt.scatter(x[:, 0], x[:, 1], c="yellow", marker='4', label='cluster7')

# u = np.array(model.u)
# plt.scatter(u[:, 0], u[:, 1], c="black", marker='X', label='center')
# plt.xlabel('petallength')
# plt.ylabel('petalwidth')
# plt.legend(loc=2)
# plt.show()
